import json     # json处理库
from datetime import datetime
import backtrader as bt
import akshare as ak
import pandas as pd

class MyStrategy(bt.Strategy):
    """
    这个就是策略配置，Backtrader
    """
    # 这个是设置20天，
    params = (("maperiod", 20),)  # 全局设定交易策略的参数

    def __init__(self):
        """
        初始化函数
        """
        self.data_close = self.datas[0].close  # 指定价格序列
        # 初始化交易指令、买卖价格和手续费
        self.order = None
        self.buy_price = None
        self.buy_comm = None
        # 添加移动均线指标
        # 这个地方是使用策略，SimplyMovingAverage普通均线策略
        # 将得到的sma(简单移动平均线)保存在变量中
        # period是前面配置的20天，就是将20天内的收盘价和除以20天
        self.sma = bt.indicators.SimpleMovingAverage(
            self.datas[0], period=self.params.maperiod
        )

    def next(self):
        """
        执行逻辑，这个也是Backtrader的
        """
        if self.order:  # 检查是否有指令等待执行,
            return
        # 检查是否持仓
        if not self.position:  # 没有持仓
            # 根据每一步的sma参数，也就是收盘价超过最近20天的收盘价平均线，就买入
            if self.data_close[0] > self.sma[0]:  # 执行买入条件判断：收盘价格上涨突破20日均线
                self.order = self.buy(size=100)  # 执行买入
        else:
            # 如果收盘价低于最近20天的收盘价平均线sma就卖出
            if self.data_close[0] < self.sma[0]:  # 执行卖出条件判断：收盘价格跌破20日均线
                self.order = self.sell(size=100)  # 执行卖出


def get_a_code_pnl(code,cur_idx,start_cash = 1000000,shouxufei=0.002):
    '''
    计算指定股票的pnl
    pnl = 最终的资金 - 初始资金
    程序推荐股票功能就是假设您携带一定资金，根据一定的策略，模拟买卖股票，得到最终手里还有多少钱

    初始资金是可以
    :param code: 股票代码
    :param cur_idx: 当前索引
    :param start_cash: 初始资金 默认1000000
    :param shouxufei: 交易手续费 默认 0.002
    :return:
    '''
    #################################################################
    # 这一串代码是获取指定股票历史数据，然后将格式转换成backtrader支持的格式
    print(f'{cur_idx}: 股票代码:{code}',end='')
    # 利用 AKShare 获取股票的后复权数据，这里只获取前 7 列
    # AKShare是一个开源的金融平台，获取数据
    stock_hfq_df = ak.stock_zh_a_hist(symbol=code, adjust="hfq").iloc[:, :7]
    # 删除 `股票代码` 列，实际运行发现，从AKShare获取的没有股票代码数据，这里的处理方式是跳过
    try:
        del stock_hfq_df['股票代码']
    except:
        return
    # 处理字段命名，以符合 Backtrader 的要求
    # Backtrader是开源的库，这些是转换成这个库的命名方式
    stock_hfq_df.columns = [
        'date',
        'open',
        'close',
        'high',
        'low',
        'volume',
    ]
    ########################################################################
    # 下面是Backtrader的标准流程
    # 把 date 作为日期索引，以符合 Backtrader 的要求
    # 分别是初始化调参运行
    stock_hfq_df.index = pd.to_datetime(stock_hfq_df['date'])
    cerebro = bt.Cerebro()  # 初始化回测系统
    start_date = datetime(2010, 4, 3)  # 回测开始时间
    end_date = datetime(2024, 6, 16)  # 回测结束时间
    data = bt.feeds.PandasData(dataname=stock_hfq_df, fromdate=start_date, todate=end_date)  # 加载数据
    cerebro.adddata(data)  # 将数据传入回测系统
    cerebro.addstrategy(MyStrategy)  # 将交易策略加载到回测系统中
    cerebro.broker.setcash(start_cash)  # 设置初始资本
    cerebro.broker.setcommission(commission=shouxufei)  # 设置交易手续费为 0.2%
    try:
        cerebro.run()  # 运行回测系统
    except Exception as e:
        print(f"{e}")
        return
    # port_value是推测出来的最终资金还剩下多少
    port_value = cerebro.broker.getvalue()  # 获取回测结束后的总资金
    # 拿最后的钱数减去一开始的本金就是pnl
    # pnl>0说明赚了多少
    # pnl<0说明亏了多少
    pnl = port_value - start_cash  # 盈亏统计
    print(f' pnl值:{pnl}')
    # 返回得到的pnl
    return pnl


def get_cur_data(data,write=True,savename=''):
    '''
    解析DataFrame，并以json的格式保存
    :param data: 数据
    :param write: 是否保存到文本
    :param savename: 保存的名称
    :return:
    '''
    # 创建JSON对象
    # 生成的json是这个格式
    '''
    {
        "data":[
            {
                "序号":"",
                "股票代码":"",
                ...
            },
            {
            },
            ...
        ]
    }
    '''
    # 这个是JSON标准流程
    # 从字符串加载，设置内容，dumps转换成json字符串，然后写入文件保存
    j = json.loads('{}')


    # 开始设置内容
    lst = []
    for one_data in data.values:
        # 存储一个股票数据
        one = {}
        for idx,value in enumerate(data.columns.tolist()):
            if value in ['上市日期']:
                # 上市日期是一个Datetime.Date类型的参数
                # 所以要转换成字符串
                # print(one_data[idx].strftime("%Y-%m-%d"))
                one[value] = one_data[idx].strftime("%Y-%m-%d")
                continue
            # 其他直接保存
            one[value] = one_data[idx]
            # 添加进列表
            lst.append(one)
    # 将列表作为value,key是data
    # 这个是json要求的格式key:value
    j['data'] = lst

    # 如果write参数是True，保存，这个是我自己设置的函数参数
    if write:
        # 转换成json的字符串,ensure_ascii是因为有中文
        j = json.dumps(j,ensure_ascii=False)
        # 写入文件
        with open(f'{savename}.json','w',encoding='utf-8') as f:
            f.write(j)

def get_all_cur_data():
    '''

    :return:
    '''
    # 通过AKShare的接口获取实时的沪深京A股数据，数据格式是pandas的DataFrame
    stock_zh_a_spot_em_df = ak.stock_zh_a_spot_em()
    # 将得到的DataFrame格式数据传入函数get_cur_data
    get_cur_data(stock_zh_a_spot_em_df, True, "实时数据沪深京A股")


if __name__ == '__main__':

    # 开始下载
    print('开始下载')
    get_all_cur_data()
    print('下载完成！')

    ##############################################
    # 下面是推荐股票的内容
    # 获取数据
    # 获取股票列表，依然是pandasd的DataFrame格式
    stock_zh_a_spot_em_df = ak.stock_zh_a_spot_em()
    # 只需要序号和代码这两列
    # 序号是为了输出好看，用于后面输出现在进行到第几个了，事实上计算只需要代码这一列，代码列的内容是股票的代码
    c = stock_zh_a_spot_em_df.loc[:,['序号','代码']]
    print('开始计算：')

    # 执行函数生成每只股票的pnl
    # 通过apply调用get_a_code_pnlh函数，并将结果保存在pnl列
    # 比如
    '''
        原本的格式是这样
        代码    序号
        代码1   0
        代码2   1
        .........全部内容
        
        执行完apply后
        代码   序号   pnl
        代码1  0     代码1通过函数get_a_code_pnl得到的pnl
        .....其他行的数据
    '''
    c['pnl'] = c.apply(lambda c:get_a_code_pnl(c['代码'],c['序号']),axis = 1)
    print("计算完成！")

    ################################################################################


    # pnl就是最终的盈亏，比如投入100块钱买股票，然后最后到手120，pnl就是120-100 = 20 赚了20
    #                    投入100块钱买股票，最后到手80      pnl就是80-100  = -20 亏了20
    # 根据pnl列值进行排序，从大到小排
    c = c.sort_values(by=['pnl'],ascending=False)
    # 排序完后，第一行元素就是pnl最大的值
    # 取前十行就是最赚的10个
    c_max_10 = c.iloc[:10]
    print(c_max_10)

    # 然后就是一样的保存流程
    # 创建json对象，赋值，转换成字符串，保存到文本文件
    j_max = json.loads('{}')
    datas = []
    for i in range(10):
        data = {}
        data['代码'] = c_max_10.iloc[i].tolist()[1]
        data['pnl'] = c_max_10.iloc[i].tolist()[2]
        datas.append(data)
    j_max['data'] = datas
    j_max = json.dumps(j_max,ensure_ascii=False)
    with open('最推荐10股','w',encoding='utf-8') as f:
        f.write(j_max)

    # 获取pnl列最大值
    # 第一行就是最赚的，所以直接输出第一行
    print(f'推荐购买股票:{c_max_10.iloc[0].tolist()[1]},预取收益:{c_max_10.iloc[0].tolist()[2]}')